import face_recognition
import cv2
import time
import psycopg2


def find_matching_face_in_db(unknown_encoding):
    """在数据库中查找最相似的人脸"""
    conn = psycopg2.connect(
        host="localhost", database="ai", user="postgres", password="liqian"
    )
    with conn.cursor() as cur:
        cur.execute(
            """
            SELECT name, 1 - (face <=> %s::vector) as similarity
            FROM face 
            WHERE 1 - (face <=> %s::vector) > 0.6
            ORDER BY face <=> %s::vector
            LIMIT 1
        """,
            (unknown_encoding.tolist(),) * 3,
        )
        return cur.fetchone()


def main():
    # 初始化摄像头
    cap = cv2.VideoCapture(0)
    if not cap.isOpened():
        print("错误：无法访问摄像头")
        return

    print(
        """
    操作指南：
    - 按 A 键：捕捉当前画面并识别人脸
    - 按 Q 键或ESC：退出程序
    - 直接关闭窗口亦可退出
    """
    )

    try:
        while True:
            ret, frame = cap.read()
            if not ret:
                print("错误：无法读取视频流")
                break

            # 显示实时画面
            cv2.imshow("人脸识别 - 按A识别/Q退出", frame)

            # 键盘输入检测（缩短等待时间提高响应速度）
            key = cv2.waitKey(50) & 0xFF

            # 退出机制（多种方式）
            if key in (27, ord("q"), ord("Q")):  # ESC或Q键退出
                print("用户主动退出程序")
                break
            elif (
                cv2.getWindowProperty("人脸识别 - 按A识别/Q退出", cv2.WND_PROP_VISIBLE)
                < 1
            ):  # 窗口关闭退出
                print("窗口已关闭")
                break

            # 人脸识别触发
            if key in (ord("a"), ord("A")):
                process_frame(frame)

    finally:
        # 确保资源释放
        cap.release()
        cv2.destroyAllWindows()
        print("程序已安全退出")


def process_frame(frame):
    """处理单帧图像的人脸识别"""
    start_time = time.time()

    try:
        # 人脸检测
        face_locations = face_recognition.face_locations(frame)
        if not face_locations:
            print("提示：未检测到人脸，请确保面部清晰可见")
            return

        # 获取编码
        unknown_encoding = face_recognition.face_encodings(frame, face_locations)[0]

        # 数据库查询
        match = find_matching_face_in_db(unknown_encoding)

        # 结果展示
        if match:
            name, similarity = match
            print(f"识别成功: {name} (相似度: {similarity:.2%})")

            # 在画面上标注结果
            top, right, bottom, left = face_locations[0]
            cv2.rectangle(frame, (left, top), (right, bottom), (0, 255, 0), 2)
            cv2.putText(
                frame,
                f"{name} {similarity:.0%}",
                (left, top - 10),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.8,
                (0, 255, 0),
                2,
            )
        else:
            print("提示：未找到匹配的人脸")
            cv2.putText(
                frame, "Unknown", (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2
            )

        # 显示结果画面（3秒）
        cv2.imshow("识别结果", frame)
        cv2.waitKey(3000)
        cv2.destroyWindow("识别结果")

    except Exception as e:
        print(f"处理出错: {str(e)}")
    finally:
        print(f"处理耗时: {time.time() - start_time:.2f}秒")


if __name__ == "__main__":
    main()
